"""
 Copyright (c) 2020 Intel Corporation
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
      http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

import copy
import io
import numpy as np
from PIL import Image
from pycocotools import coco
from pycocotools import mask as mask_api

from beta.examples.tensorflow.common.object_detection.utils import box_utils
from beta.examples.tensorflow.common.object_detection.utils import mask_utils


class COCOWrapper(coco.COCO):
    """COCO wrapper class.

    This class wraps COCO API object, which provides the following additional functionalities:
        1. Support string type image id.
        2. Support loading the groundtruth dataset using the external annotation dictionary.
        3. Support loading the prediction results using the external annotation dictionary.
    """

    def __init__(self, eval_type='box', annotation_file=None, gt_dataset=None):
        """Instantiates a COCO-style API object.

        Args:
            eval_type: either 'box' or 'mask'.
            annotation_file: a JSON file that stores annotations of the eval dataset.
                This is required if `gt_dataset` is not provided.
            gt_dataset: the groundtruth eval datatset in COCO API format.
        """
        if ((annotation_file and gt_dataset) or ((not annotation_file) and (not gt_dataset))):
            raise ValueError('One and only one of `annotation_file` and `gt_dataset` needs to be specified.')

        if eval_type not in ['box', 'mask']:
            raise ValueError('The `eval_type` can only be either `box` or `mask`.')

        coco.COCO.__init__(self, annotation_file=annotation_file)
        self._eval_type = eval_type
        if gt_dataset:
            self.dataset = gt_dataset
            self.createIndex()

    def load_res(self, predictions):
        """Loads result file and return a result api object.

        Args:
            predictions: a list of dictionary each representing an annotation in COCO
                format. The required fields are `image_id`, `category_id`, `score`,
                `bbox`, `segmentation`.

        Returns:
            res: result COCO api object.

        Raises:
            ValueError: if the set of image id from predctions is not the subset of
                the set of image id of the groundtruth dataset.
        """
        res = coco.COCO()
        res.dataset['images'] = copy.deepcopy(self.dataset['images'])
        res.dataset['categories'] = copy.deepcopy(self.dataset['categories'])

        image_ids = [ann['image_id'] for ann in predictions]
        if set(image_ids) != (set(image_ids) & set(self.getImgIds())):
            raise ValueError('Results do not correspond to the current dataset!')
        for ann in predictions:
            x1, x2, y1, y2 = [ann['bbox'][0], ann['bbox'][0] + ann['bbox'][2],
                              ann['bbox'][1], ann['bbox'][1] + ann['bbox'][3]]
            if self._eval_type == 'box':
                ann['area'] = ann['bbox'][2] * ann['bbox'][3]
                ann['segmentation'] = [[x1, y1, x1, y2, x2, y2, x2, y1]]
            elif self._eval_type == 'mask':
                ann['area'] = mask_api.area(ann['segmentation'])

        res.dataset['annotations'] = copy.deepcopy(predictions)
        res.createIndex()

        return res


def convert_predictions_to_coco_annotations(predictions):
    """Converts a batch of predictions to annotations in COCO format.

    Args:
      predictions: a dictionary of lists of numpy arrays including the following
          fields. K below denotes the maximum number of instances per image.
          Required fields:
            - source_id: a list of numpy arrays of int or string of shape
                [batch_size].
            - num_detections: a list of numpy arrays of int of shape [batch_size].
            - detection_boxes: a list of numpy arrays of float of shape
                [batch_size, K, 4], where coordinates are in the original image
                space (not the scaled image space).
            - detection_classes: a list of numpy arrays of int of shape
                [batch_size, K].
            - detection_scores: a list of numpy arrays of float of shape
                [batch_size, K].
          Optional fields:
            - detection_masks: a list of numpy arrays of float of shape
                [batch_size, K, mask_height, mask_width].

    Returns:
      coco_predictions: prediction in COCO annotation format.
    """
    coco_predictions = []
    num_batches = len(predictions['source_id'])
    batch_size = predictions['source_id'][0].shape[0]
    max_num_detections = predictions['detection_classes'][0].shape[1]
    use_outer_box = 'detection_outer_boxes' in predictions
    for i in range(num_batches):
        predictions['detection_boxes'][i] = box_utils.yxyx_to_xywh(predictions['detection_boxes'][i])
        if use_outer_box:
            predictions['detection_outer_boxes'][i] = box_utils.yxyx_to_xywh(predictions['detection_outer_boxes'][i])
            mask_boxes = predictions['detection_outer_boxes']
        else:
            mask_boxes = predictions['detection_boxes']

        for j in range(batch_size):
            if 'detection_masks' in predictions:
                image_masks = mask_utils.paste_instance_masks(
                    predictions['detection_masks'][i][j],
                    mask_boxes[i][j],
                    int(predictions['image_info'][i][j, 0, 0]),
                    int(predictions['image_info'][i][j, 0, 1]))
                binary_masks = (image_masks > 0.0).astype(np.uint8)
                encoded_masks = [
                    mask_api.encode(np.asfortranarray(binary_mask))
                    for binary_mask in list(binary_masks)]

            for k in range(max_num_detections):
                ann = {}
                ann['image_id'] = predictions['source_id'][i][j]
                ann['category_id'] = predictions['detection_classes'][i][j, k]
                ann['bbox'] = predictions['detection_boxes'][i][j, k]
                ann['score'] = predictions['detection_scores'][i][j, k]

                if 'detection_masks' in predictions:
                    ann['segmentation'] = encoded_masks[k]

                coco_predictions.append(ann)

    for i, ann in enumerate(coco_predictions):
        ann['id'] = i + 1

    return coco_predictions


def convert_groundtruths_to_coco_dataset(groundtruths, label_map=None):
    """Converts groundtruths to the dataset in COCO format.

    Args:
      groundtruths: a dictionary of numpy arrays including the fields below.
        Note that each element in the list represent the number for a single
        example without batch dimension. K below denotes the actual number of
        instances for each image.
        Required fields:
          - source_id: a list of numpy arrays of int or string of shape
            [batch_size].
          - height: a list of numpy arrays of int of shape [batch_size].
          - width: a list of numpy arrays of int of shape [batch_size].
          - num_detections: a list of numpy arrays of int of shape [batch_size].
          - boxes: a list of numpy arrays of float of shape [batch_size, K, 4],
              where coordinates are in the original image space (not the
              normalized coordinates).
          - classes: a list of numpy arrays of int of shape [batch_size, K].
        Optional fields:
          - is_crowds: a list of numpy arrays of int of shape [batch_size, K]. If
              th field is absent, it is assumed that this instance is not crowd.
          - areas: a list of numy arrays of float of shape [batch_size, K]. If the
              field is absent, the area is calculated using either boxes or
              masks depending on which one is available.
          - masks: a list of numpy arrays of string of shape [batch_size, K],
      label_map: (optional) a dictionary that defines items from the category id
        to the category name. If `None`, collect the category mappping from the
        `groundtruths`.

    Returns:
      coco_groundtruths: the groundtruth dataset in COCO format.
    """
    source_ids = np.concatenate(groundtruths['source_id'], axis=0)
    heights = np.concatenate(groundtruths['height'], axis=0)
    widths = np.concatenate(groundtruths['width'], axis=0)
    gt_images = [{'id': int(i), 'height': int(h), 'width': int(w)} for i, h, w
                in zip(source_ids, heights, widths)]

    gt_annotations = []
    num_batches = len(groundtruths['source_id'])
    batch_size = groundtruths['source_id'][0].shape[0]
    for i in range(num_batches):
        for j in range(batch_size):
            num_instances = groundtruths['num_detections'][i][j]
            for k in range(num_instances):
                ann = {}
                ann['image_id'] = int(groundtruths['source_id'][i][j])
                if 'is_crowds' in groundtruths:
                    ann['iscrowd'] = int(groundtruths['is_crowds'][i][j, k])
                else:
                    ann['iscrowd'] = 0
                ann['category_id'] = int(groundtruths['classes'][i][j, k])
                boxes = groundtruths['boxes'][i]
                ann['bbox'] = [float(boxes[j, k, 1]),
                               float(boxes[j, k, 0]),
                               float(boxes[j, k, 3] - boxes[j, k, 1]),
                               float(boxes[j, k, 2] - boxes[j, k, 0])]
                if 'areas' in groundtruths:
                    ann['area'] = float(groundtruths['areas'][i][j, k])
                else:
                    ann['area'] = float((boxes[j, k, 3] - boxes[j, k, 1]) * (boxes[j, k, 2] - boxes[j, k, 0]))

                if 'masks' in groundtruths:
                    mask = Image.open(io.BytesIO(groundtruths['masks'][i][j, k]))
                    width, height = mask.size
                    np_mask = (np.array(mask.getdata()).reshape(height, width).astype(np.uint8))
                    np_mask[np_mask > 0] = 255
                    encoded_mask = mask_api.encode(np.asfortranarray(np_mask))
                    ann['segmentation'] = encoded_mask
                    if 'areas' not in groundtruths:
                        ann['area'] = mask_api.area(encoded_mask)

                gt_annotations.append(ann)

    for i, ann in enumerate(gt_annotations):
        ann['id'] = i + 1

    if label_map:
        gt_categories = [{'id': i, 'name': label_map[i]} for i in label_map]
    else:
        category_ids = [gt['category_id'] for gt in gt_annotations]
        gt_categories = [{'id': i} for i in set(category_ids)]

    gt_dataset = {
        'images': gt_images,
        'categories': gt_categories,
        'annotations': copy.deepcopy(gt_annotations),
    }

    return gt_dataset
